import os
import torch
import numpy as np
from dataLoad import DataLoader


class Calculate:
    def __init__(self):
        """
        :param vec -> tensor
        :param word -> ndarray
        """
        self.dataLoader = DataLoader(r'./data/glove.6B.300d.txt')
        self.vec = torch.from_numpy(self.dataLoader.get_vec().astype(np.float64))
        self.word = self.dataLoader.get_word()
        self.cosine_path = os.path.abspath('./cos_sim')
        self.topk_path = r'./toplist/topk.pth'
        self.topk_nd = None

    @staticmethod
    def __preprocess_data(self):
        """
        send data to gpu for calculating the cos-sim
        :param self:
        :return:
        """
        gpu = torch.device('cuda:0')
        self.vec = self.vec.to(gpu)

    @staticmethod
    def __cosine_distance(E, rawNorm, start=0, step=500):
        """
        :param E:              vec of glove
        :param rawNorm:        Calculate the cosine similarity of how many rows
        :param start:          "the start of each batch"
        :param step:           "batch size"
        :return:               cos-sim matrix
        """
        dot = E[start: start + step, :] @ E.t()
        norm = torch.norm(E[start: start + step, :], 2, 1)

        # print(dot.shape)
        # print(norm.shape)

        x = torch.div(dot, torch.unsqueeze(rawNorm, 0))  # 500 * 40w / 1 * 40w     div 2-norm of other vec
        x = torch.div(x, torch.unsqueeze(norm, 1))  # 1w * 40w / 500 * 4      div 2-norm of vec itself

        return x

    def cosine_all_to_file(self, totalrows=400000, useGPU=True):
        """
        :param totalrows:  Choose how many rows of cos_sim to calculate
        :param useGPU:     decide whether to use gpu
        :return:           nothing but save result to local file
        """
        if useGPU is True:
            Calculate.__preprocess_data(self)
        else:
            pass

        raw_norm = torch.norm(self.vec, 2, 1)
        # TODO : torch save
        for i in range(int(totalrows / 500)):
            res = Calculate.__cosine_distance(self.vec, raw_norm, start=i * 500, step=500)
            torch.save(res, os.path.join(self.cosine_path, '{}_cos.pth'.format(i * 500)))
            # np.save(r'./cos_res/{}_cos'.format(i * 500), res.cpu().numpy())

    def sim2word(self, src: str, dest: str):
        """
        :param src:   src word
        :param dest:  dest word
        :return:      the cosine similarity of src and dest
        """
        # word in the glove is unique
        idx1 = np.where(self.word == src)[0][0]  # [0][0] convert array to int
        idx2 = np.where(self.word == dest)[0][0]

        idx = int(idx1 / 500) * 500
        # file = os.path.join(self.cosine_path, '{}_cos.pth'.format(idx))
        file = r'./cos_res/{}_cos.npy'.format(idx)
        # cosine_sim = torch.load(file)
        cos_sim = np.load(file)

        return cos_sim[idx1 % 500][idx2]

    def cal_top_k(self, K=100, rows=400000):
        """
        save
        :param K:    top k sim word
        :return:     idx of top-k words
        """
        topk = []
        # TODO torch read
        for i in range(int(rows / 500)):
            name = '{}_cos.pth'.format(i * 500)
            file = os.path.join(self.cosine_path, name)
            # data = np.load(file)
            data = torch.load(file)
            data_1 = data[:50].to("cuda:0")
            data_2 = data[50:100].to("cuda:1")
            data_3 = data[100:150].to("cuda:2")
            data_4 = data[150:200].to("cuda:3")
            idx_1 = torch.argsort(data_1, dim=1)[:, -1 * K - 1:-1]
            torch.save(idx_1, self.topk_path)
            idx_2 = torch.argsort(data_2, dim=1)[:, -1 * K - 1:-1]
            idx_3 = torch.argsort(data_3, dim=1)[:, -1 * K - 1:-1]
            idx_4 = torch.argsort(data_4, dim=1)[:, -1 * K - 1:-1]
            topk.append(idx_1)
            topk.append(idx_2)
            topk.append(idx_3)
            topk.append(idx_4)

        self.topk_nd = topk

    def load_topk(self):
        """
        Show the words with the greatest cosine similarity of the first five words
        :return:
        """
        # path = os.path.join(self.topk_path, 'topk.pth')

        data = torch.load(self.topk_path)
        res = []
        for i in range(5):
            word = self.word[data[i].cpu()]
            res.append(word)

        print(self.word[[0, 1, 2, 3, 4]])
        res = np.array(res)
        np.savetxt('./top5-top100.txt', res, fmt='%s')
        return res

    def inquire_for_topk(self, src: str):
        idx = np.where(self.word == src)[0][0]
        top_k_word = self.word[self.topk_nd[idx]]

        return top_k_word


if __name__ == '__main__':
    calc = Calculate()
    print("father, mother", calc.sim2word('father', 'mother'))
    print("king, queen", calc.sim2word('king', 'queen'))
    print("man, woman", calc.sim2word('man', 'woman'))
